"""
MQTT客户端模块 (改进版)
提供易于使用的MQTT功能接口，支持TLS加密、自动重连，以及按主题分发消息回调。
"""

import paho.mqtt.client as mqtt
import logging
from typing import Optional, Callable, Any, Dict, Tuple
from dataclasses import dataclass
from enum import Enum

logger = logging.getLogger(__name__)

# ... (QoSLevel 和 MQTTConfig 类保持不变) ...
class QoSLevel(Enum):
    AT_MOST_ONCE = 0
    AT_LEAST_ONCE = 1
    EXACTLY_ONCE = 2

@dataclass
class MQTTConfig:
    host: str = "ta9fe3f1.ala.cn-hangzhou.emqxsl.cn"
    port: int = 8883
    username: Optional[str] = None
    password: Optional[str] = None
    client_id: Optional[str] = None
    keepalive: int = 60
    clean_session: bool = True
    automatic_reconnect: bool = True
    min_retry_interval: int = 1
    max_retry_interval: int = 60
    enable_tls: bool = True
    tls_ca_cert: Optional[str] = "emqxsl-ca (1).crt"
    tls_certfile: Optional[str] = None
    tls_keyfile: Optional[str] = None
    tls_insecure: bool = False


class MQTTClient:
    """MQTT客户端类 (改进版)
    
    提供完整的MQTT客户端功能，包括连接管理、消息发布/订阅、TLS加密、
    自动重连，并支持为不同主题指定不同的消息回调函数。
    """
    
    def __init__(self, config: MQTTConfig):
        self.config = config
        self.client = mqtt.Client(
            client_id=config.client_id,
            clean_session=config.clean_session
        )
        
        if config.username and config.password:
            self.client.username_pw_set(config.username, config.password)
        
        if config.enable_tls:
            self._configure_tls()
        
        # === MODIFIED: 内部回调和数据结构的改变 ===
        self.client.on_connect = self._on_connect
        self.client.on_disconnect = self._on_disconnect
        self.client.on_publish = self._on_publish
        # 设置内部统一的消息入口
        self.client.on_message = self._on_message_dispatcher 

        # 用于存储所有订阅信息（主题 -> (QoS, 特定回调函数))
        self._subscriptions: Dict[str, Tuple[int, Optional[Callable]]] = {}
        # 用于存储默认的回调函数
        self._default_on_message: Optional[Callable] = None

    # ... (_configure_tls, _on_publish, _on_disconnect 方法保持不变) ...
    def _configure_tls(self):
        try:
            if self.config.tls_ca_cert:
                self.client.tls_set(
                    ca_certs=self.config.tls_ca_cert,
                    certfile=self.config.tls_certfile,
                    keyfile=self.config.tls_keyfile
                )
            else:
                self.client.tls_set()
            if self.config.tls_insecure:
                self.client.tls_insecure_set(True)
        except Exception as e:
            logging.error(f"TLS配置失败: {e}")
            raise

    def _on_publish(self, client, userdata, mid):
        logger.debug(f"消息发布成功，消息ID: {mid}")
        
    def _on_disconnect(self, client, userdata, rc):
        if rc != 0:
            logger.warning(f"MQTT客户端意外断开连接，错误码: {rc}")
        else:
            logger.info("MQTT客户端已正常断开连接")
        if self.config.automatic_reconnect:
            logger.info("自动重连已启用，将自动尝试重新连接")

    def _on_connect(self, client, userdata, flags, rc):
        if rc == 0:
            logger.info(f"成功连接到MQTT服务器 {self.config.host}:{self.config.port}")
            
            # === MODIFIED: 重连后重新订阅 ===
            # 重新订阅时，不需要再次绑定回调，因为它们已经存储在 self._subscriptions 中
            if self._subscriptions:
                logger.info("重新订阅之前订阅的主题")
                for topic, (qos, _) in self._subscriptions.items():
                    self.client.subscribe(topic, qos)
                    logger.info(f"重新订阅主题: {topic} (QoS: {qos})")
        else:
            logger.error(f"连接MQTT服务器失败，错误码: {rc}")

    # ... (connect 和 disconnect 方法保持不变) ...
    def connect(self) -> bool:
        try:
            if self.config.automatic_reconnect:
                self.client.reconnect_delay_set(
                    self.config.min_retry_interval,
                    self.config.max_retry_interval
                )
            self.client.connect(
                self.config.host,
                self.config.port,
                self.config.keepalive
            )
            self.client.loop_start()
            return True
        except Exception as e:
            logger.error(f"连接MQTT服务器时出错: {e}")
            return False

    def disconnect(self):
        self.client.loop_stop()
        self.client.disconnect()

    def publish(self, topic: str, payload: Any, qos: QoSLevel = QoSLevel.AT_MOST_ONCE, retain: bool = False) -> bool:
        try:
            result = self.client.publish(topic, payload, qos.value, retain)
            if result.rc == mqtt.MQTT_ERR_SUCCESS:
                logger.debug(f"消息已发布到主题: {topic}")
                return True
            else:
                logger.error(f"发布消息失败，错误码: {result.rc}")
                return False
        except Exception as e:
            logger.error(f"发布消息时出错: {e}")
            return False
            
    def subscribe(self, topic: str, qos: QoSLevel = QoSLevel.AT_MOST_ONCE, callback: Optional[Callable] = None):
        """订阅指定主题，并可以选择性地为其绑定一个特定的回调函数。
        
        Args:
            topic: 要订阅的主题，可以包含通配符(+, #)。
            qos: QoS等级。
            callback: (可选) 一个专门处理此主题消息的回调函数。
        """
        # === MODIFIED: 存储主题、QoS和特定的回调 ===
        logger.info(f"订阅主题 '{topic}'，并绑定回调: {callback.__name__ if callback else '无'}")
        self._subscriptions[topic] = (qos.value, callback)
        self.client.subscribe(topic, qos.value)

    # === NEW: 消息分发器 ===
    def _on_message_dispatcher(self, client, userdata, msg):
        """
        内部消息分发器。
        当收到消息时，此方法会查找最匹配的回调函数并执行它。
        它支持通配符，并有一个默认回调作为备用。
        """
        logger.debug(f"收到消息，开始分发 - 主题: {msg.topic}")
        
        # 尝试解码消息内容
        try:
            payload_str = msg.payload.decode()
        except UnicodeDecodeError:
            payload_str = msg.payload.decode('utf-8', errors='replace')
            logger.warning(f"消息内容解码时出现错误，已使用替换字符处理: {msg.topic}")
        logger.debug(f"内容: {payload_str}")

        handled = False
        # 遍历所有订阅，查找匹配的回调
        for sub_topic, (_, callback) in self._subscriptions.items():
            if mqtt.topic_matches_sub(sub_topic, msg.topic):
                if callback:
                    try:
                        callback(client, userdata, msg)
                        handled = True
                    except Exception as e:
                        logger.error(f"执行主题 '{msg.topic}' 的回调函数时出错: {e}")

        # 如果没有找到任何特定的回调函数，则调用默认回调
        if not handled and self._default_on_message:
            logger.debug(f"主题 '{msg.topic}' 没有找到特定回调，使用默认回调。")
            try:
                self._default_on_message(client, userdata, msg)
            except Exception as e:
                logger.error(f"执行默认回调函数时出错: {e}")

    # === NEW: 设置默认回调的方法 ===
    def set_default_message_callback(self, callback: Callable):
        """设置一个默认的消息回调函数。
        
        当收到的消息没有匹配到任何通过 subscribe() 指定的特定回调时，
        此默认回调函数将被调用。

        Args:
            callback: 消息回调函数，签名应为 callback(client, userdata, msg)
        """
        self._default_on_message = callback
        
    @property
    def is_connected(self) -> bool:
        return self.client.is_connected()